from construct import *
from construct.text import *



#===============================================================================
# AST transformations
#===============================================================================
class NumberNode(Adapter):
    def _decode(self, obj, context):
        sign, integral, frac, denom = obj
        if not sign:
            sign = "+"
        if frac:
            base, exp = frac
            if exp:
                expsign, expval = exp
                if not expsign:
                    expsign = "+"
                exptext = "e" + sign + expval
            else:
                exptext = ""
            value = float(sign + integral + "." + base + exptext)
        else:
            value = int(sign + integral)
        if denom:
            denom = int("".join(denom))
            value = float(value) / denom
        return AstNode("num", value = value)

class StringNode(Adapter):
    def _decode(self, obj, context):
        return AstNode("str", value = obj)

class ListNode(Adapter):
    def _decode(self, obj, context):
        elems, lastelem = obj
        if lastelem:
            elems.append(lastelem)
        return AstNode("list", value = elems)

class SymNode(Adapter):
    def _decode(self, obj, context):
        return AstNode("sym", name = obj)

class CallNode(Adapter):
    def _decode(self, obj, context):
        expr, args, lastarg = obj
        if lastarg:
            args.append(lastarg)
        return AstNode("call", func = expr, args = args)

class CommentNode(Adapter):
    def _decode(self, obj, context):
        return AstNode("comment", text = obj)

class LambdaNode(Adapter):
    def _decode(self, obj, context):
        args, lastarg, func = obj
        if lastarg:
            args.append(lastarg)
        return AstNode("lambda", args = [a.name for a in args], func = func)

class DefStmtNode(Adapter):
    def _decode(self, obj, context):
        name, arglist, lastarg, func = obj
        if lastarg:
            arglist.append(lastarg)
        args = [a.name for a in arglist]
        return AstNode("def", name = name.name, args = args, func = func)

class EvalStmtNode(Adapter):
    def _decode(self, obj, context):
        return AstNode("eval", func = obj)


#===============================================================================
# macros
#===============================================================================
def OptSeq(name, *args):
    return Optional(Sequence(name, *args))
def SeqOfOne(name, *args):
    return IndexingAdapter(Sequence(name, *args), index = 0)
def OptSeqOfOne(name, *args):
    return Optional(SeqOfOne(name, *args))

ws = Whitespace(" \t\r\n")
rws = Whitespace(" \t\r\n", optional = False)
digits = StringAdapter(GreedyRange(CharOf("digit", "0123456789")))
expr = LazyBound("expr", lambda: expr)


#===============================================================================
# grammar rules
#===============================================================================
num_node = NumberNode(
        Sequence("num",
        Optional(CharOf("sign", "-")),
        digits,
        OptSeq("fraction",
            Literal("."),
            digits,
            OptSeq("exp",
                Literal("e"),
                Optional(CharOf("sign", "+-")),
                digits,
            ),
        ),
        OptSeqOfOne("denom",
            Literal("/"),
            digits,
        )
    )
)

str_node = StringNode(
    QuotedString("str", encoding = "string-escape")
)

list_node = ListNode(
    Sequence("list",
        Literal("["),
        OptionalGreedyRange(
            SeqOfOne("elem",
                expr,
                Literal(","),
            )
        ),
        Optional(expr),
        Optional(Literal(",")),
        Literal("]"),
    )
)

sym_node = SymNode(
    Identifier("sym")
)

call_node = CallNode(
    Sequence("call",
        Select("expr",
            SeqOfOne("paren",
                Literal("("),
                expr,
                Literal(")"),
            ),
            sym_node,
        ),
        ws,
        Literal("("),
        OptionalGreedyRange(
            SeqOfOne("elem",
                expr,
                Literal(","),
            )
        ),
        Optional(expr),
        Optional(Literal(",")),
        ws,
        Literal(")"),
    )
)

lambda_node = LambdaNode(
    Sequence("lambda",
        Literal("<"),
        ws,
        OptionalGreedyRange(
            SeqOfOne("args",
                sym_node,
                ws,
                Literal(","),
                ws,
            )
        ),
        Optional(sym_node),
        ws,
        Literal(":"),
        expr,
        Literal(">"),
    )
)

comment_node = CommentNode(
    SeqOfOne("comment",
        Literal("#"),
        Line("text"),
    )
)

expr = SeqOfOne("expr",
    ws,
    Select("expr",
        comment_node,
        num_node,
        str_node,
        list_node,
        lambda_node,
        call_node,
        sym_node,
    ),
    ws,
)

def_stmt = DefStmtNode(
    Sequence("def",
        ws,
        sym_node,
        rws,
        OptionalGreedyRange(
            SeqOfOne("arglist", 
                sym_node, 
                ws,
                Literal(","),
                ws,
            )
        ),
        ws,
        Optional(sym_node),
        ws,
        Literal("="),
        expr,
    )
)

eval_stmt = EvalStmtNode(
    expr,
)

root = OptionalGreedyRange(
    Select("root",
        def_stmt,
        eval_stmt,
        comment_node,
    )
)






